package com.fwpsl.graph.minimumspanningtree;

import com.fwpsl.graph.Edge;
import com.fwpsl.graph.Graph;
import com.fwpsl.graph.Vertex;

import java.util.*;

/**
 * @author: 风戏fw
 * @date: 2024/4/5
 * @description: Prim算法是一种用于求解图的最小生成树的算法。该算法的基本思想是从一个起始节点开始，选择与当前生成树相连的边中权值
 * 最小的边，然后将该边加入到生成树中，再选择下一条权值最小的边，直到生成树覆盖了所有节点为止。
 * Prim算法是一种贪心算法，通过每一步选择当前情况下最优的边来逐步构建最小生成树。其时间复杂度为O(V^2)或O(E*logV)，
 * 取决于采用何种数据结构来实现。
 * <p>
 * 在实际应用中，Prim算法常用于网络设计、电缆布线、电路设计等需要在节点之间建立连接的问题，
 * 能够以最小的总成本连接所有节点，从而节省资源和成本。
 */
public class Prim {

    static class PathNode<T> {
        Vertex<T> currentVertex;
        Vertex<T> preVertex;
        double distance;

        public PathNode(Vertex<T> currentVertex) {
            this.currentVertex = currentVertex;
            this.preVertex = null;
            this.distance = Double.MAX_VALUE;
        }
    }

    public static <T> void prim(Graph<T> graph) {
        // visited数组来记录哪些节点已经被加入了最小生成树
        Map<Vertex<T>, Boolean> visitedMap = new HashMap<>();

        // dist数组来记录每个节点与最小生成树中已有节点的最短距离；pre数组来记录与当前节点最短距离的节点是哪个
        Map<Vertex<T>, PathNode<T>> distances = new HashMap<>();
        for (Vertex<T> vertex : graph.getVertices()) {
            distances.put(vertex, new PathNode<>(vertex));
        }
        // 将第一个顶点的关键点值设为0，以便首先提取它
        Vertex<T> vertex = graph.getFirstVertex();
        distances.get(vertex).distance = 0;

        PriorityQueue<Edge> pq = new PriorityQueue<>(Comparator.comparingDouble(Edge::getWeight));
        pq.add(new Edge(vertex, vertex, 0));
        while (!pq.isEmpty()) {
            Edge u = pq.poll();

            // 如果MST中有顶点，则继续下一个边
            if (Boolean.TRUE.equals(visitedMap.get(u.getDestination()))) {
                continue;
            }
            visitedMap.put(u.getDestination(), true);

            // 更新拾取顶点的相邻顶点的关键值，并将它们添加到最小堆中
            for (Edge<T> edge : ((Vertex<T>) u.getDestination()).getEdges()) {
                PathNode<T> pathNode = distances.get(edge.getDestination());
                // 拓展未到达的顶点，并且是最短的路径的情况。
                if (!Boolean.TRUE.equals(visitedMap.get(edge.getDestination())) && edge.getWeight() < pathNode.distance) {
                    pathNode.preVertex = u.getDestination();
                    pathNode.distance = edge.getWeight();
                    pq.add(edge);
                }
            }
        }

        // 打印构建的MST
        System.out.println("Edge \tWeight");
        for (Vertex<T> v : graph.getVertices()) {
            PathNode<T> pathNode = distances.get(v);
            if (pathNode.preVertex == null) {
                continue;
            }
            System.out.println(pathNode.preVertex.getData() + " - " + v.getData() + "\t" + pathNode.distance);
        }
    }

    public static void main(String[] args) {
        int[][] graphEdgeArr = {
                {0, 4, 0, 0, 0, 0, 0, 9, 0},
                {4, 0, 8, 0, 0, 0, 0, 11, 0},
                {0, 8, 0, 7, 0, 4, 0, 0, 2},
                {0, 0, 7, 0, 9, 14, 0, 0, 0},
                {0, 0, 0, 9, 0, 10, 0, 0, 0},
                {0, 0, 4, 14, 10, 0, 2, 0, 0},
                {0, 0, 0, 0, 0, 2, 0, 1, 6},
                {9, 11, 0, 0, 0, 0, 1, 0, 7},
                {0, 0, 2, 0, 0, 0, 6, 7, 0}
        };
        Graph<Integer> graph = Graph.buildIntGraph(graphEdgeArr);

        prim(graph);
    }
}
